import paramiko
import threading
from threading import Thread
import traceback
import socket
import json
from io import StringIO
from paramiko.ssh_exception import NoValidConnectionsError
from paramiko.ssh_exception import AuthenticationException
from paramiko.ssh_exception import SSHException

from utils.custom_log   import log_start
logger = log_start('ssh')


zmodemszstart = b'rz\r**\x18B00000000000000\r\x8a'
zmodemszend = b'**\x18B0800000000022d\r\x8a'
zmodemrzstart = b'rz waiting to receive.**\x18B0100000023be50\r\x8a'
zmodemrzend = b'**\x18B0800000000022d\r\x8a'
zmodemcancel = b'\x18\x18\x18\x18\x18\x08\x08\x08\x08\x08'


class SSH():
    def __init__(self,conndict, websocker, message):
        self.host = conndict["host"]
        self.username = conndict["user"]
        self.port = conndict["port"]
        self.password = conndict.get("password", "")
        self.private = conndict.get("private", "")
        self.width = conndict["width"]
        self.height = conndict["height"]
        self.websocker = websocker
        self.message = message
        self.timeout = 60
        self.connect_()

    def connect_(self):
        try:
            self.sshclient = paramiko.SSHClient()
            self.sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            if self.private :
                private = StringIO(self.private)
                key = paramiko.RSAKey.from_private_key(private)
                self.sshclient.connect(hostname=self.host, port=self.port, username=self.username, pkey=key,
                                       timeout=self.timeout,allow_agent=False, look_for_keys=False)
                logger.info('sshkey认证登录')
                # look_for_keys: 是否在~/.ssh中搜索私钥文件
                # allow_agent: 是否允许连接到ssh代理
            else:
                self.sshclient.connect(hostname=self.host, port=self.port, username=self.username, password=self.password,
                timeout=self.timeout,allow_agent=False, look_for_keys=False)
                logger.info('密码认证登录')

            transport = self.sshclient.get_transport()
            self.channel = transport.open_session()
            self.channel.get_pty(term='xterm', width=self.width, height=self.height)
            self.channel.invoke_shell()

            for i in range(2):
                recv = self.channel.recv(1024).decode('utf-8')
                self.message['status'] = 200
                self.message['message'] = recv
                message = json.dumps(self.message)
                self.websocker.send(message)

        except NoValidConnectionsError as connect:
            msg = {'status': 1, 'message': "无法建立SSH连接：%s，请重试！" % connect}
            logger.error(msg)
            self.websocker.send(json.dumps(msg))
            self.websocker.close()

        except AuthenticationException as auth:
            msg = {'status': 1, 'message': '验证错误: %s,检查用户密码!!! ' % auth}
            logger.error(msg)
            self.websocker.send(json.dumps(msg))
            self.websocker.close()

        except SSHException as ssh:
            msg = {'status': 1, 'message': "SSH连接失败: %s,请检查重试" % ssh}
            logger.error(msg)
            self.websocker.send(json.dumps(msg))
            self.websocker.close()

        except Exception as  e:
            msg = {'status': 1, 'message': "%s"  % e}
            logger.error(msg)
            self.websocker.send(json.dumps(msg))
            self.websocker.close()

    def resize_pty(self, cols, rows):
        self.channel.resize_pty(width=cols, height=rows)

    def django_to_ssh(self, data):
        try:
            # print(data.decode('utf-8'))
            # if data == 'exit':
            #     self.channel.send(data)
            # else:
            #     self.channel.send(data)
            #     self.close_ssh()
            self.channel.send(data)
    
        except Exception as  e:
            self.close_ssh('连接SSH异常,关闭连接,%s' % e)

    def websocket_to_django(self):
        try:
            while True:
                data = self.channel.recv(1024).decode('utf-8', 'ignore')
                if not len(data):
                    return
                self.message['status'] = 0
                self.message['message'] = data
                message = json.dumps(self.message)
                self.websocker.send(message)
        except Exception as  e:
            self.close_ssh('发送到WSS异常,关闭连接,%s' % e)

    def close_ssh(self,message='关闭连接'):
        self.message['status'] = 1
        self.message['message'] = message
        message = json.dumps(self.message)
        self.websocker.send(message)
        self.websocker.close()
        self.channel.close()

    def task_shell(self, data):
        Thread(target=self.django_to_ssh, args=(data,)).start()
        Thread(target=self.websocket_to_django).start()